import json
import random
from collections import defaultdict
import os

def split_dataset(json_path, train_ratio=0.7):
    # 读取JSON文件
    with open(json_path, 'r') as f:
        data = json.load(f)

    # 按类别分组图片
    category_images = defaultdict(list)
    for image in data['images']:
        file_name = image['file_name']
        # 提取类别（第一个下划线前的部分）
        if '_' in file_name:
            category = file_name.split('_', 1)[0]
            category_images[category].append(image)

    # 计算总图片数和类别数
    total_images = sum(len(imgs) for imgs in category_images.values())
    total_categories = len(category_images)
    print(f"总类别数: {total_categories}, 总图片数: {total_images}")

    # 将类别随机排序
    categories = list(category_images.keys())
    random.shuffle(categories)

    # 初始化训练集和验证集
    train_images = []
    val_images = []
    train_count = 0
    val_count = 0
    train_categories = []
    val_categories = []

    # 计算目标训练集大小
    target_train_size = int(total_images * train_ratio)

    # 分配类别到训练集或验证集
    for category in categories:
        images = category_images[category]
        num_images = len(images)

        # 如果当前训练集图片数 + 该类图片数 <= 目标大小，则加入训练集
        if train_count + num_images <= target_train_size:
            train_images.extend(images)
            train_count += num_images
            train_categories.append(category)
        else:
            # 如果加入训练集会超出目标大小，则加入验证集
            val_images.extend(images)
            val_count += num_images
            val_categories.append(category)

    # 计算最终比例
    actual_ratio = train_count / total_images
    print(f"\n训练集: {len(train_images)} 张图片 ({len(train_categories)} 个类别)")
    print(f"验证集: {len(val_images)} 张图片 ({len(val_categories)} 个类别)")
    print(f"实际分割比例: 训练集 {actual_ratio:.2%}, 验证集 {1 - actual_ratio:.2%}")

    # 为每个图像添加split字段
    for img in train_images:
        img['split'] = 'train'
    for img in val_images:
        img['split'] = 'val'

    # 返回划分结果
    return {
        "train": train_images,
        "val": val_images,
        "train_categories": train_categories,
        "val_categories": val_categories
    }


def check_image_matching(json_path, folder_path, extra_output_txt="extra_images.txt"):
    # 读取JSON文件
    with open(json_path, 'r') as f:
        data = json.load(f)

    # 获取所有JSON中的文件名
    json_files = set()
    for image in data['images']:
        file_name = image['file_name']
        # 确保文件名是bmp格式（不区分大小写）
        if not file_name.lower().endswith('.bmp'):
            file_name += '.bmp'
        json_files.add(file_name.lower())  # 统一转换为小写比较

    # 获取文件夹中的所有BMP图片
    folder_files = set()
    # 同时保存原始文件名（用于输出）
    original_folder_files = []
    for file in os.listdir(folder_path):
        if file.lower().endswith('.bmp'):
            folder_files.add(file.lower())
            original_folder_files.append(file)  # 保存原始文件名

    # 检查缺失的图片（JSON中有但文件夹中没有）
    missing_in_folder = json_files - folder_files
    # 检查多余的图片（文件夹中有但JSON中没有）
    extra_in_folder = folder_files - json_files

    # 输出结果
    print(f"JSON记录数量: {len(json_files)} | 文件夹图片数量: {len(folder_files)}")
    print("\n=== 缺失的图片（在JSON中但不在文件夹中）===")
    print('\n'.join(missing_in_folder) if missing_in_folder else "无缺失图片")

    print(f"\n=== 多余的图片（在文件夹中但不在JSON中）: {len(extra_in_folder)} 张 ===")
    print('\n'.join(extra_in_folder) if extra_in_folder else "无多余图片")

    # 将多余图片的名称（不带后缀）写入txt文件
    if extra_in_folder:
        # 创建映射：小写文件名 -> 原始文件名
        lower_to_original = {f.lower(): f for f in original_folder_files}

        # 获取不带后缀的原始文件名
        extra_names_no_ext = []
        for lower_name in extra_in_folder:
            original_name = lower_to_original.get(lower_name, lower_name)
            # 移除.bmp后缀
            if original_name.lower().endswith('.bmp'):
                name_no_ext = original_name[:-4]  # 移除.bmp
            else:
                name_no_ext = original_name
            extra_names_no_ext.append(name_no_ext)

        # 写入文件
        with open(extra_output_txt, 'w') as f:
            for name in extra_names_no_ext:
                f.write(name + '\n')
        print(f"\n已将 {len(extra_names_no_ext)} 个多余图片名称写入: {extra_output_txt}")
    else:
        print("无多余图片，未创建输出文件")


# 使用示例
if __name__ == "__main__":
    # json_file = r"D:\Code\0-data\1-齿轮检测数据集\青山数据集\qingshan_dianquketi_v1.json"  # 替换为你的JSON文件路径
    #
    # # 执行划分
    # result = split_dataset(json_file)
    #
    # # 可选：保存划分后的JSON
    # # with open("split_dataset.json", "w") as outfile:
    # #     json.dump(result, outfile, indent=2)
    #
    # # 可选：打印类别划分详情
    # print("\n=== 训练集类别 ===")
    # print(", ".join(sorted(result["train_categories"])))
    #
    # print("\n=== 验证集类别 ===")
    # print(", ".join(sorted(result["val_categories"])))

    json_file = r"D:\Code\0-data\1-齿轮检测数据集\青山数据集\qingshan_dianquketi_v1.json"  # 替换为你的JSON文件路径
    images_folder = r"D:\Code\0-data\1-齿轮检测数据集\青山数据集\VOCdevkit\VOC2012\JPEGImages"  # 替换为你的图片文件夹路径
    check_image_matching(json_file, images_folder,
                         r'D:\Code\0-data\1-齿轮检测数据集\青山数据集\VOCdevkit\VOC2012\ImageSets\Main\extra_images.txt')